'use strict'

const fs = require('node:fs')

const fsp = fs.promises

/**
 * 基于Linux cgroups v2 的资源限制。
 * 
 * 任何操作都要在/sys/fs/cgroup中进行
 * 
 * 创建资源组：
 * 
 * mkdir cf
 * 
 * cgroup.controllers	表示可以使用的资源控制，根cgroup内包含全部可用的资源控制，而子cgroup可以使用的控制继承自其父组的cgroup。
 * cgroup.subtree_control	表示当前已经启用的资源控制，其内容会继承到子cgroup的cgroup.controllers中。
 * 
 * 因为subtree_control的继承关系，需要确保在/sys/fs/cgroup/cgroup.subtree_control中必须要存在子组需要的项。
 * 
 * #启用对cpu memory的资源控制，+ 表示启用 - 表示不启用
 * echo '+cpu +memory +pids' > cgroup.subtree_control
 * 
 * 启用了pids，就可以把进程放在 cgroup.procs中进行管理
 * 
 * echo 1234 > cgroup.procs
 * echo 1256 >> cgroup.procs
 * 
 * 这表示CPU的10000个时间片，占用5000个，就是最大50%占有率。
 * echo 5000 10000 > cpu.max
 *
 * 对内存的限制，还需要设置memory.swap.max 和 memory.zswap.max
 * 否则内存超出限制会写入swap而不是终止进程。
 * 
 * 对内存的限制是整个组的限制，而不是单个进程。CPU的限制也是如此。
 * CPU会让整个cgroup中的进程分配总共的占比，这进行过严格测试。
 * 但是对内存的测试并不严格。
 */
//'domain threaded', 'domain invalid'
//以上在Linux 内核6.4.12上测试无效
//threaded模式开启后，不可以再变回domain
//threaded模式是针对线程的，也就是说如果cpu限制在

//domain模式会让每个线程均分CPU资源，而采用threaded模式就是每个线程占用CPU资源。

let allowType = [
  'domain', 'threaded'
]

class CGroup {

  constructor() {

    this.cgdir = '/sys/fs/cgroup'

    /**
     * 完整的信息请参考：
     * @link https://docs.kernel.org/admin-guide/cgroup-v2.html?highlight=admin+guide+cgroup+v2+rst
     * rbps  Max read bytes per second
     * wbps Max write bytes per second
     * riops Max read IO operations per second
     * wiops Max write IO operations per second
     * IO限制的格式：
     * 8:0 rbps=2097152 wbps=max riops=max wiops=120
     * 如果你去限制一个分区会报告没有这个设备，这个错误信息提示不准确。
     * 这种原因是因为linux只能对设备限制IO，不能对分区限制。
     * 从 /proc/partions获取分区信息，并得到设备的主设备号以及对应的表示设备的从设备号，这里是0。
     * 注意不要使用分区。
     * {
     *    [NAME] : {
     *      dirname: '',
     *      control: [
     *        'cpu', 'memory', 'pids'
     *      ],
     *      cpu: 50 | [3000, 10000],
     *      memory: 100000
     *    }
     * }
     */
    this.groupTable = {}

    this.defaultControl = ['cpuset', 'cpu', 'memory', 'pids', 'io']

    Object.defineProperty(this, 'allowType', {
      configurable: false,
      writable: false,
      enumerable: true,
      value: allowType
    })

    this.namePreg = /^[a-z][a-z0-9_\-]{1,100}$/i

    this.cpuset = {
      effective: null
    }

    this.getEffectiveCPU().then(res => {
      if (res.ok) {
        this.cpuset.effective = res.cpus
      }
    })
  }

  parseCPUS(data) {
    let cpus = []
    let cpumap = {}

    data.split(',').filter(p => p.length > 0).forEach(x => {
      if (x.indexOf('-') > 0) {
        let arr = x.split('-')
        let start = parseInt(arr[0])
        let end = parseInt(arr[1])
        for (let i = start; i <= end; i++) {
          if (!cpumap[`cpu_${i}`]) {
            cpumap[`cpu_${i}`] = true
            cpus.push(i)
          }
        }
      } else {
        if (!cpumap[`cpu_${x}`]) {
          cpumap[`cpu_${x}`] = true
          cpus.push(parseInt(x))
        }
      }
    })

    return cpus
  }

  async getEffectiveCPU(grp = '', cpufile='cpuset.cpus.effective') {
    try {
      let cpus_eft = this.cgdir + (grp ? '/' + grp : '') + '/' + cpufile
      await fsp.access(cpus_eft)
      //格式：0-3,5,6-9,12
      let data = await fsp.readFile(cpus_eft, {encoding: 'utf8'})
    
      return {
        ok: true,
        cpus: this.parseCPUS(data)
      }
    } catch (err) {
      return {
        ok: false,
        error: err
      }
    }
  }

  checkDetail(detail) {
    if (!detail.control || !Array.isArray(detail.control) || detail.control.length === 0) {
      if (!detail.control || typeof detail.control !== 'string') {
        detail.control = [...this.defaultControl]
      }
    }

    /**
     * 默认是domain、可以是 threaded、domain threaded、domain invalid
     * threaded模式在没有任何任务运行后，会变回domain模式，但是cgroup.type仍然是threaded
     */
    if (!detail.type || this.allowType.indexOf(detail.type) < 0) {
      detail.type = 'domain'
    }

    //maj min rbps wbps riops wiops
    if (!detail.io) detail.io = ''
    else if (typeof detail.io !== 'object') {
      detail.io = {}
    }
    
    if (detail.pids === 0 || detail.pids === 'max') {
      detail.pids = 'max'
    }
    else if (!detail.pids || typeof detail.pids !== 'number' || detail.pids < 0) {
      detail.pids = 'max'
    }

    //字节为单位 
    if (detail.memory !== 'max' 
        && (!detail.memory || typeof detail.memory !== 'number' || detail.memory < 1000)
    ) {
      detail.memory = 'max'
    }

    if (detail.swap === undefined) {
      detail.__swap__ = false
    } else {
      detail.__swap__ = true
    }

    if (detail.swap !== 'max' && (typeof detail.swap !== 'number' || detail.swap < 0)) {
      if (detail.memory === 'max')
        detail.swap = 'max'
      else detail.swap = 0
    }

    if (!detail.cpu || (!Array.isArray(detail.cpu) && (typeof detail.cpu !== 'number' || detail.cpu < 0))) {
      detail.cpu = [9500, 10000]
    } else if (typeof detail.cpu === 'number') {
      if (detail.cpu >= 100) detail.cpu = 0
      else {
        let a = detail.cpu
        detail.cpu = [a * 100, 10000]
      }
    }

    //有效CPU范围，解析出来后从
    detail.cpuset = {
      effective: null
    }

  }

  parseCPULimit(detail) {
    if (!detail.cpus || !detail.cpuset || !detail.cpuset.effective) return ''

    let cpus = detail.cpus

    let elength = detail.cpuset.effective.length
    let ecpus = detail.cpuset.effective

    if (elength === 1) return ''

    if (cpus[0] === '%') {
      let cp, pos = 'rand'

      if (['-', '+', '='].indexOf(cpus[cpus.length - 1]) >= 0) {
        cp = parseInt(cpus.substring(1, cpus.length - 1))
        pos = cpus[cpus.length - 1]
      } else {
        cp = parseInt(cpus.substring(1))
      }
      
      if (isNaN(cp)) return ''

      if (cp <= 0 || cp >= 100) return ''

      cp = cp / 100

      let ind = 0
      switch (pos) {
        case 'rand':
          ind = parseInt(Math.random() * (elength - elength * cp))
          return ecpus.slice(ind, parseInt(elength * cp) + ind)
          break
        
        case '=':
          ind = parseInt(elength / 2) - parseInt(elength * cp / 2)
          return ecpus.slice(ind, ind + parseInt(elength * cp + 0.5))
          break

        case '-':
          ind = 0
          return ecpus.slice(ind, parseInt(elength * cp + 0.2))
          break

        case '+':
          let total = parseInt(elength * cp + 0.1)
          let real_cpus = []
          let count = 0
          for (let i = elength-1; i >= 0; i--) {
            real_cpus.push(ecpus[i])
            count++
            if (count >= total) break
          }

          return real_cpus.sort((a,b) => {
            return a-b
          })
      }
      
    }

    return cpus
  }

  async setType(name, typ) {
    if (this.allowType.indexOf(typ) < 0) return false
    
    let grp_path = this.findPath(name)
    if (!grp_path) return false

    return await fsp.writeFile(grp_path + '/cgroup.type', typ)
                    .then(r => {
                      this.groupTable[name].type = typ
                      return true
                    })
                    .catch(err => {
                      return false
                    })
  }

  async create(name, detail) {
    if (!this.namePreg.test(name)) {
      throw new Error(`${name} 不符合要求，要求：字母开头，支持字母、数字、下划线、连字符（减号），长度2～100。`)
    }

    this.checkDetail(detail)

    this.groupTable[name] = detail

    let cgrp_subtree_control = this.cgdir + '/cgroup.subtree_control'
  
    let context = await fsp.readFile(cgrp_subtree_control, {encoding: 'utf8'})
  
    let carr = context.split(' ').filter(p => p.length > 0)
  
    let add_controllers = []
    for (let x of this.defaultControl) {
      if (carr.indexOf(x) < 0) {
        add_controllers.push('+' + x)
      }
    }

    if (add_controllers.length > 0) {
      await fsp.writeFile(cgrp_subtree_control, add_controllers.join(' '), {flag: 'a'})
    }
    
    let cgrpdir = this.cgdir + '/' + name

    try {
      await fsp.access(cgrpdir)
    } catch (err) {
      await fsp.mkdir(cgrpdir)
    }

    try {
      let cur_type = await fsp.readFile(cgrpdir + '/cgroup.type', {encoding: 'utf8'})
      if (cur_type !== 'threaded' && cur_type !== detail.type) {
        await this.setType(name, detail.type)
      }
    } catch (err) {}

    if (Array.isArray(detail.cpu)) {
      await fsp.writeFile(cgrpdir + '/cpu.max', `${detail.cpu[0]} ${detail.cpu[1]}`)
    }

    if (detail.memory && detail.type === 'domain') {
      await fsp.writeFile(cgrpdir + '/memory.max', `${detail.memory}`).catch(err => {
        console.error(err)
      })

      if (detail.memory !== 'max') {
        await this.setSwap(name, detail.__swap__ ? detail.swap : 0)
      }
    }

    await this.setIO(name, detail.io)
    await this.setPids(name, detail.pids)

    if (detail.control.indexOf('cpuset') >= 0) {
      let cpu_eft = await this.getEffectiveCPU(detail.name)
      if (cpu_eft.ok) {
        detail.cpuset.effective = cpu_eft.cpus
      }
    }

    /**
     * 尽可能动态支持：
     *  百分比自动确定：%50 %80, %50- %50+ %20=
     *      - 表示前50%，+表示后50% =表示从中间开始
     *      默认是随机处理。
     *  指定：0-1,2
     */
    if (detail.cpus && typeof detail.cpus !== 'string') {
      detail.cpus = ''
    }

    if (detail.cpus) {
      await this.setCPUSet(name, detail.cpus)
    }
  }

  findPath(name) {
    let cgrp = this.groupTable[name]
    if (!cgrp) return false
    return this.cgdir + '/' + name
  }

  //不必自己删除，让cdpc调用rmdir -rf 删除即可。
  remove(name) {
    let cgrp = this.groupTable[name]
    if (!cgrp) return false
  }

  fmtIO(iobj) {
    if (iobj.maj === undefined) return ''

    let iostr_arr = []

    ;['rbps', 'wbps', 'riops', 'wiops'].forEach(x => {
      if (iobj[x]) {
        iostr_arr.push(`${x}=${iobj[x]}`)
      }
    })

    if (iostr_arr.length === 0) return ''

    return `${iobj.maj}:${iobj.min || 0} ${iostr_arr.join(' ')}`
  }

  async setCPUSet(name, cpus) {
    let cgrp = this.groupTable[name]
    if (!cgrp) return false

    cgrp.__cpus__ = cpus
    let real_cpus = this.parseCPULimit(cgrp)

    cgrp.cpus = Array.isArray(real_cpus) ? real_cpus.join(',') : real_cpus

    if (cgrp.cpus) {
      try {
        let cpusetfile = this.cgdir + '/' + name + '/cpuset.cpus'
        await fsp.access(cpusetfile)
        await fsp.writeFile(cpusetfile, cgrp.cpus, {encoding:'utf8'})
      } catch (err) {
        return false
      }
    }

    return true
  }

  async setPids(name, pid_max) {
    let cgrp = this.groupTable[name]
    if (!cgrp) return false

    if (!pid_max || (pid_max !== 'max' && typeof pid_max !== 'number') ) {
      return false
    }

    try {
      await fsp.writeFile(this.cgdir + '/' + name + '/pids.max', `${pid_max}`)
      cgrp.pids = pid_max
      return true
    } catch (err) {
      console.error(err)
      return false
    }

  }

  async setIO(name, iobj) {
    if (!iobj || typeof iobj !== 'object') return false
    
    let cgrp = this.groupTable[name]
    if (!cgrp) return false

    let iostr = this.fmtIO(iobj)
    if (!iostr) return false

    try {
      await fsp.writeFile(this.cgdir + '/' + name + '/io.max', iostr)
      cgrp.io = iobj
      return true
    } catch (err) {
      console.error(err)
      return false
    }
  }

  async setCPU(name, cpu) {
    if (!cpu && typeof cpu !== 'number') return false

    if (typeof cpu === 'number') {
      if (cpu < 0) return false

      if (cpu === 0 || cpu >= 100) {
        cpu = [10000, 10000]
      }
    }

    let cgrp = this.groupTable[name]
    if (!cgrp) return false

    //detail.cpu = cpu

    try {
      await fsp.writeFile(this.cgdir + '/' + name + '/cpu.max', `${cpu[0]} ${cpu[1]}`)
      detail.cpu = cpu
      return true
    } catch (err) {
      return false
    }
  }

  async setSwap(name, mem) {
    let cgrp = this.groupTable[name]
    if (!cgrp) return false

    try {
      await fsp.writeFile(this.cgdir + '/' + name + '/memory.swap.max', `${cgrp.swap}`)
      cgrp.swap = mem
      let zswap_file = this.cgdir + '/' + name + '/memory.zswap.max'
      await fsp.access(zswap_file)
        .then(async res => {
          return await fsp.writeFile(zswap_file, `${cgrp.swap}`)
        })
        .catch(err => {})

      return true
    } catch (err) {
      console.error(err)
      return false
    }
  }

  async setMem(name, mem, setSwap=true) {
    if (!mem && typeof mem !== 'number') return false
    
    if (mem <= 0) mem = 'max'

    let cgrp = this.groupTable[name]
    if (!cgrp) return false

    //线程模式，没有memory.max文件
    if (cgrp.type !== 'domain') return false

    if (mem !== 'max' && setSwap) {
      await this.setSwap(name, 0)
    }

    try {
      await fsp.writeFile(this.cgdir + '/' + name + '/memory.max', `${cgrp.memory}`)
      cgrp.memory = mem
      return true
    } catch (err) {
      return false
    }
  }

  async addPids(name, pids) {
    if (typeof pids === 'number') pids = [pids]

    if (!Array.isArray(pids)) return false

    let realPids = []

    pids.forEach(p => {
      if (typeof p === 'number' && p > 1) {
        realPids.push(p)
      }
    })

    if (realPids.length <= 0) return false

    let cgrp = this.groupTable[name]
    if (!cgrp) return false

    try {
      await fsp.writeFile(this.cgdir + '/' + name + '/cgroup.procs', realPids.join('\n'), {flag: 'a'})
    } catch (err) {
      return false
    }

    return true
  }

}

module.exports = CGroup
